from typing import List, Optional, Mapping, Any
from functools import partial

from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
import urllib.request
from urllib import parse


### chatglm-6B llm ###
class ChatGLM(LLM):
    max_length: int = 2048
    temperature: float = 0.95
    top_p: float = 0.7
    model_path: str = ""
    history: List = []
    streaming: bool = True

    def __init__(self, top_p=0.7, temperature=0.95, max_length=2048, **kwargs: Any):
        super().__init__(**kwargs)
        self.top_p = top_p
        self.temperature = temperature
        self.max_length = max_length
        self.streaming = True
        self.model_path = ""
        self.history = []


    @property
    def _llm_type(self) -> str:
        return "chatglm-6B"

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            "max_length": self.max_length,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "history": [],
            "streaming": self.streaming
        }

    def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            add_history: bool = False  # 是否将此次对话加入到历史列表中
    ) -> str:

        if self.streaming:
            text_callback = partial(StreamingStdOutCallbackHandler().on_llm_new_token, verbose=True)
            resp = self.generate_resp(prompt, text_callback, add_history=add_history)
        else:
            resp = self.generate_resp(self, prompt, add_history=add_history)

        return resp

    def generate_resp(self, prompt, text_callback=None, add_history=True):
        resp = ""
        index = 0
        dict = {
            'prompt': prompt

        }
        data = bytes(parse.urlencode(dict), encoding="utf8")
        url = "http://172.16.4.242:7860/generate_respe"
        # url = "http://42.48.115.228:62001/generate_respe"
        response = urllib.request.urlopen(url, data=data)
        resp = response.read().decode("utf-8")
        if text_callback:

            if add_history:
                self.history += [[prompt, resp]]
            text_callback(resp)
        else:
            if add_history:
                self.history += [[prompt, resp]]
        return resp

    def set_params(self, **kwargs):
        for k, v in kwargs.items():
            if k in self._identifying_params:
                self.k = v

# model = ChatGLM()
# resp_1 = model._call("你好")
# print(resp_1)
# model = ChatGLM(0.95)
